import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os
import glob
from scipy import stats

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 安全文件读取函数
def safe_read_file(file_pattern, skiprows=None, **kwargs):
    """尝试多种方式查找并读取文件"""
    matching_files = glob.glob(file_pattern)
    if not matching_files:
        raise FileNotFoundError(f"找不到匹配的文件: {file_pattern}")
    
    file_path = matching_files[0]
    print(f"找到文件: {file_path}")
    
    encodings = ['utf-8', 'gbk', 'latin1', 'ISO-8859-1', 'cp1252']
    for encoding in encodings:
        try:
            if file_path.endswith('.csv'):
                df = pd.read_csv(file_path, encoding=encoding, skiprows=skiprows, **kwargs)
            elif file_path.endswith(('.xls', '.xlsx')):
                df = pd.read_excel(file_path, skiprows=skiprows, **kwargs)
            else:
                raise ValueError(f"不支持的文件格式: {file_path}")
            print(f"成功读取: {file_path} (编码: {encoding})")
            return df
        except UnicodeDecodeError:
            continue
        except Exception as e:
            print(f"读取文件时出错: {file_path}, 错误: {str(e)}")
            continue
            
    raise UnicodeDecodeError(f"所有编码方式均失败: {file_path}")

def load_index_data(index_name, earliest_date):
    """加载指数数据并应用最早日期过滤"""
    try:
        file_pattern = os.path.join(current_dir, f'*{index_name}*指数*.csv')
        df = safe_read_file(file_pattern, skiprows=1, header=None)
        
        # 检查列数并设置列名
        if len(df.columns) >= 5:
            df.columns = ['date', 'open', 'high', 'low', 'close'][:len(df.columns)]
        else:
            df = df.iloc[:, :2]
            df.columns = ['date', 'close']
        
        # 日期处理
        df['date'] = pd.to_datetime(df['date'], errors='coerce')
        df = df.dropna(subset=['date'])
        
        # 应用最早日期过滤
        if not df.empty:
            df = df[df['date'] >= pd.to_datetime(earliest_date)]
            print(f"{index_name}指数数据加载成功，原始{len(df)}行，过滤后{len(df)}行")
        else:
            print(f"{index_name}指数数据加载成功，但数据为空")
            
        return df[['date', 'close']]
    except Exception as e:
        print(f"加载{index_name}指数数据时出错: {str(e)}")
        return pd.DataFrame(columns=['date', 'close'])

def main():
    # 获取当前工作目录
    current_dir = os.getcwd()
    print(f"当前工作目录: {current_dir}")
    
    # 定义各指数最早日期
    index_earliest_dates = {
        'IM': '2010-04-16',
        'IC': '2015-04-16',
        'IF': '2010-04-16',  # 假设IF最早日期与IM相同
        'IH': '2015-04-16'
    }

    # 加载净空单数据
    try:
        net_short = safe_read_file(
            os.path.join(current_dir, '*net_short*.csv'),
            skiprows=1, 
            header=None
        )
        # 处理列名
        if len(net_short.columns) >= 6:
            net_short.columns = ['date', 'IF', 'IC', 'IM', 'IH', 'total'][:len(net_short.columns)]
        else:
            # 如果列数不足，使用前几列
            net_short.columns = ['date', 'IF', 'IC', 'IM', 'IH'][:len(net_short.columns)]
        
        # 日期处理
        net_short['date'] = pd.to_datetime(net_short['date'], format='%Y%m%d', errors='coerce')
        net_short = net_short.dropna(subset=['date'])
        print(f"净空单数据加载成功，共{len(net_short)}行")
        print(net_short.head())
    except Exception as e:
        print(f"加载净空单数据时出错: {str(e)}")
        return

    # 加载各指数数据
    im_data = load_index_data('IM', index_earliest_dates['IM'])
    ic_data = load_index_data('IC', index_earliest_dates['IC'])
    if_data = load_index_data('IF', index_earliest_dates['IF'])
    ih_data = load_index_data('IH', index_earliest_dates['IH'])

    # 合并数据，确保以净空单日期为准
    all_data = net_short.copy()

    # 先对指数数据按日期排序，确保合并时顺序正确
    for name, df in [('IM', im_data), ('IC', ic_data), ('IF', if_data), ('IH', ih_data)]:
        if not df.empty:
            df = df.sort_values('date')  # 按日期排序
            all_data = pd.merge(all_data, df, on='date', how='left', suffixes=('', f'_{name}'))
            all_data = all_data.rename(columns={'close': f'{name}_close'})
        else:
            print(f"警告: {name}指数数据为空，跳过合并")

    # 再次按日期排序，确保数据顺序一致
    if 'date' in all_data.columns:
        all_data = all_data.sort_values('date')

    # 确保有必要的列
    required_columns = ['date', 'IF', 'IC', 'IM', 'IH']
    for col in required_columns:
        if col not in all_data.columns:
            print(f"警告: 缺少必要列 {col}，使用0填充")
            all_data[col] = 0

    # 计算净空单增量
    for col in ['IF', 'IC', 'IM', 'IH']:
        if col in all_data.columns:
            all_data[f'{col}_Delta'] = all_data[col].diff()
        else:
            print(f"警告: 无法计算 {col}_Delta，列不存在")
            all_data[f'{col}_Delta'] = np.nan

    # 计算涨跌幅
    for idx in ['IM', 'IC', 'IF', 'IH']:
        close_col = f'{idx}_close'
        if close_col in all_data.columns:
            all_data[f'{idx}_Return'] = all_data[close_col].pct_change() * 100
        else:
            print(f"警告: 无法计算 {idx}_Return，列不存在")
            all_data[f'{idx}_Return'] = np.nan

    # 计算T+1, T+2, T+3的涨跌幅
    for idx in ['IM', 'IC', 'IF', 'IH']:
        return_col = f'{idx}_Return'
        if return_col in all_data.columns:
            for lag in [1, 2, 3]:
                all_data[f'{idx}_T{lag}'] = all_data[return_col].shift(-lag)
        else:
            print(f"警告: 无法计算 {idx}_T*，列不存在")
            for lag in [1, 2, 3]:
                all_data[f'{idx}_T{lag}'] = np.nan

    # 分析函数
    def analyze_relationship(data, index_name):
        # 检查必要的列
        delta_col = f'{index_name}_Delta'
        t1_col = f'{index_name}_T1'
        t2_col = f'{index_name}_T2'
        t3_col = f'{index_name}_T3'
        
        required_cols = [delta_col, t1_col, t2_col, t3_col]
        if not all(col in data.columns for col in required_cols):
            print(f"警告: {index_name} 缺少必要列，跳过分析")
            return None, None
        
        # 过滤有效数据
        df = data[['date', delta_col, t1_col, t2_col, t3_col]].copy()
        df = df.dropna()
        
        if df.empty:
            print(f"警告: {index_name} 没有有效数据")
            return None, None
            
        # 分组分析
        increase = df[df[delta_col] > 0]
        decrease = df[df[delta_col] < 0]
        
        # 计算平均涨跌幅
        results = {
            'Delta > 0': {
                '样本数': len(increase),
                'T+1平均涨跌': increase[t1_col].mean() if not increase.empty else np.nan,
                'T+2平均涨跌': increase[t2_col].mean() if not increase.empty else np.nan,
                'T+3平均涨跌': increase[t3_col].mean() if not increase.empty else np.nan
            },
            'Delta < 0': {
                '样本数': len(decrease),
                'T+1平均涨跌': decrease[t1_col].mean() if not decrease.empty else np.nan,
                'T+2平均涨跌': decrease[t2_col].mean() if not decrease.empty else np.nan,
                'T+3平均涨跌': decrease[t3_col].mean() if not decrease.empty else np.nan
            }
        }
        
        # 计算相关系数
        corr_t1 = df[[delta_col, t1_col]].corr().iloc[0, 1]
        
        return results, corr_t1

    # 执行分析
    results = {}
    correlations = {}

    for idx in ['IM', 'IC', 'IF', 'IH']:
        res, corr = analyze_relationship(all_data, idx)
        if res is not None:
            results[idx] = res
            correlations[idx] = corr

    # 打印结果
    if results:
        for idx, res in results.items():
            print(f"\n{idx}指数分析结果:")
            print("| 净空单变化 | 样本数 | T+1平均涨跌 | T+2平均涨跌 | T+3平均涨跌 |")
            print("|------------|--------|--------------|--------------|--------------|")
            for group, values in res.items():
                print(f"| {group} | {values['样本数']} | {values['T+1平均涨跌']:.2f}% | {values['T+2平均涨跌']:.2f}% | {values['T+3平均涨跌']:.2f}% |")
            
            print(f"\nT+1涨跌与净空单增量的相关系数: {correlations[idx]:.4f}")
    else:
        print("\n警告: 没有可用的分析结果")

    # 可视化结果
    def plot_results(index_name, results):
        if index_name not in results:
            print(f"警告: 没有 {index_name} 的结果用于绘图")
            return
            
        increase = results[index_name]['Delta > 0']
        decrease = results[index_name]['Delta < 0']
        
        fig, ax = plt.subplots(figsize=(10, 6))
        periods = ['T+1', 'T+2', 'T+3']
        
        # 净空单增加
        ax.plot(periods, [increase['T+1平均涨跌'], increase['T+2平均涨跌'], increase['T+3平均涨跌']], 
                'o-', label='净空单增加', color='red')
        
        # 净空单减少
        ax.plot(periods, [decrease['T+1平均涨跌'], decrease['T+2平均涨跌'], decrease['T+3平均涨跌']], 
                'o-', label='净空单减少', color='green')
        
        ax.set_title(f'{index_name}指数: 净空单变化对后续涨跌幅的影响')
        ax.set_ylabel('平均涨跌幅(%)')
        ax.axhline(0, color='gray', linestyle='--')
        ax.legend()
        plt.show()

    # 为每个指数生成图表
    for idx in ['IM', 'IC', 'IF', 'IH']:
        if idx in results:
            plot_results(idx, results)

    # 极端情况分析 (以IC指数为例)
    try:
        if 'IC_Delta' in all_data.columns and 'IC_T1' in all_data.columns:
            ic_data = all_data[['date', 'IC_Delta', 'IC_T1']].copy().dropna()
            extreme_increase = ic_data[ic_data['IC_Delta'] > 10000]
            
            if not extreme_increase.empty:
                down_days = len(extreme_increase[extreme_increase['IC_T1'] < 0])
                probability = down_days / len(extreme_increase) * 100

                print(f"\nIC指数极端情况分析(净空单增加>10,000手):")
                print(f"样本数量: {len(extreme_increase)}")
                print(f"T+1下跌天数: {down_days}")
                print(f"下跌概率: {probability:.2f}%")
            else:
                print("\nIC指数没有净空单增加>10,000手的样本")
        else:
            print("\n缺少IC指数数据，无法进行极端情况分析")
    except Exception as e:
        print(f"极端情况分析出错: {str(e)}")

    # 整体相关性分析
    corr_data = []
    for idx in ['IM', 'IC', 'IF', 'IH']:
        delta_col = f'{idx}_Delta'
        t1_col = f'{idx}_T1'
        
        if delta_col in all_data.columns and t1_col in all_data.columns:
            df = all_data[[delta_col, t1_col]].dropna()
            if not df.empty:
                corr = df.corr().iloc[0, 1]
                corr_data.append({'Index': idx, 'Correlation': corr})
    
    if corr_data:
        corr_df = pd.DataFrame(corr_data)
        print("\n各指数T+1涨跌与净空单增量的相关性:")
        print(corr_df)
    else:
        print("\n没有可用的相关性数据")

    # 绘制相关性热力图
    try:
        # 选择可用的列
        available_cols = []
        for idx in ['IM', 'IC', 'IF', 'IH']:
            if f'{idx}_Delta' in all_data.columns and f'{idx}_T1' in all_data.columns:
                available_cols.extend([f'{idx}_Delta', f'{idx}_T1'])
        
        if available_cols:
            corr_matrix = all_data[available_cols].corr()
            
            plt.figure(figsize=(12, 8))
            sns.heatmap(corr_matrix, annot=True, fmt=".2f", cmap='coolwarm', center=0,
                        annot_kws={"size": 10}, linewidths=0.5)
            plt.title('净空单增量与T+1涨跌幅相关性热力图')
            plt.tight_layout()
            plt.show()
        else:
            print("没有足够的数据绘制热力图")
    except Exception as e:
        print(f"绘制热力图时出错: {str(e)}")

if __name__ == "__main__":
    main()